import ctypes
import os
import sys
import cv2
from ctypes import *
from timeit import default_timer as timer

import numpy as np
from PyQt5.Qt import QImageReader
import colorsys
from PIL import Image
import configparser
from VOCdefects.Detector_Struct import *
from universal.root_path import dir_root
conf = configparser.ConfigParser()


class run:
    def __init__(self,trainWeight,img_path,class_names,imageSize,score,iou,test_goodImg_outdir,test_defectsImg_outdir,iniFilePath):
        self.iniFilePath = os.path.join(os.path.dirname(os.path.dirname(trainWeight)), iniFilePath)
        print(self.iniFilePath)
        hsv_tuples = [(x / len(class_names), 1., 1.)
                      for x in range(len(class_names))]

        self.dirpath = dir_root()
        self.dllPath = os.sep+"lib"+os.sep
        self.dllPath2 = os.sep+"TrainData"+os.sep
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        self.Objdll = ctypes.windll.LoadLibrary(self.dirpath+self.dllPath+"Detector_API.dll")
        self.boxcomplex = bboxcomplex()
        self.strcfg = bytes(self.dirpath+self.dllPath2 + "yolov4-tiny.cfg", 'utf-8')
        print(self.strcfg,'self.strcfg')
        self.strWeight = bytes(trainWeight, 'utf-8')
        # with open(dirpath + "labels_ch.txt",'w') as f:
        #     for a in class_names:
        #         f.write(a+'\n')
        # self.strLabel = bytes(dirpath + "labels_ch.txt", 'utf-8')
        # print(self.strLabel,'self.strLabel')
        self.iCount = c_int(0)
        self.Objdll.Object_Detector.restype = c_bool
        self.img_path = img_path
        self.class_names = class_names
        self.testParameters = testParameters()
        self.testParameters.classes = len(class_names)
        self.testParameters.width = int(imageSize.split('*')[0])
        self.testParameters.height = int(imageSize.split('*')[1])
        self.DefectClasses = class_names

        self.imageSize = imageSize
        self.strWeight = bytes(trainWeight, 'utf-8')

        self.testParameters.batch = 1
        self.testParameters.subdivisions =1
        self.testParameters.filters = (self.testParameters.classes + 5) * 3
        self.Objdll.ModleInitialEx.restype = c_void_p
        self.pModel = self.Objdll.ModleInitialEx(self.testParameters, self.strcfg, self.strWeight, c_int(0))
        self.test_defectsImg_outdir = test_defectsImg_outdir
        self.test_goodImg_outdir = test_goodImg_outdir
        self.SocreIou = float(score)
        self.isImgAnchored = True
        self.totalClassesDict = {}
        self.class_names.append('OK')
        self.class_names.append('NG')
        for i in self.class_names:
            self.totalClassesDict[i] = []

    def runYolo(self):
        test_images = self.scanAllImages(self.img_path)
        timeall = 0
        NGImagesNum = 0
        for img in test_images:
            beginTime = timer()
            boxes, imgClasses = self.TestImgs(img)
            endTime = timer()
            t = endTime-beginTime
            timeall = timeall+t
            size_img = Image.open(img)
            detect_img = cv2.imread(img)
            detected_img = self.detect_image(size_img,detect_img, boxes)
            img_name = os.path.basename(img)
            for box in boxes:
                if box[5] in imgClasses:
                    if img_name not in self.totalClassesDict[box[-1]]:
                        self.totalClassesDict[box[-1]].append(img_name)
                        cv2.imwrite(os.path.join(self.test_defectsImg_outdir, img_name), detected_img)
                    if img_name not in self.totalClassesDict['NG']:
                        self.totalClassesDict['NG'].append(img_name)
            if boxes == []:
                self.totalClassesDict['OK'].append(img_name)
                cv2.imwrite(os.path.join(self.test_goodImg_outdir, img_name), detected_img)
        totalTestImgsNum = len(test_images)
        NGImagesNum = len(self.totalClassesDict['NG'])

        print("测试总图像：%d,检测出的缺陷图片：%d" % (totalTestImgsNum, NGImagesNum))
        if totalTestImgsNum >0:
            ngRate = '%.4f' % (float(NGImagesNum) / totalTestImgsNum)
            NGrate = '%.2f' % (100 * float(ngRate))
        else:
            ngRate = '0.0%'
        OKImagesNum = len(self.totalClassesDict['OK'])

        totalTestTime = '%.3f' % (float(timeall))
        averageTime = '%.3f' % (float(totalTestTime) / float(totalTestImgsNum))
        # self.Objdll.UnLoadModel.restype = c_bool
        # self.Objdll.UnLoadModel(ctypes.c_void_p(self.pModel))
        return str(totalTestImgsNum), str(NGImagesNum), str(OKImagesNum), str(NGrate), str(totalTestTime), str(averageTime), self.totalClassesDict

    def TestImgs(self, imgPath):
        imgpth = bytes(imgPath, 'utf-8')
        image = cv2.imdecode(np.fromfile(imgpth,dtype=np.uint8),1)
        inputImageData = InputImage()
        inputImageData.h,inputImageData.w,inputImageData.c, = image.shape
        inputImageData.data = image.ctypes.data_as(ctypes.c_char_p)
        self.Objdll.Object_Detector(inputImageData,ctypes.c_void_p(self.pModel),
                                    ctypes.byref(self.boxcomplex),ctypes.byref(self.iCount))
        boxes = []
        imgclasses = []
        for i in range(self.iCount.value):
            prob = self.boxcomplex.box[i].prob
            prob = '%.2f' % (100 * float(prob)) + '%'
            # imgclasses.append(self.boxcomplex.box[i].labelName.decode())
            # labelname = self.boxcomplex.box[i].labelName.decode()
            labelname = self.DefectClasses[self.boxcomplex.box[i].obj_id]
            imgclasses.append(labelname)

            box = [self.boxcomplex.box[i].x, self.boxcomplex.box[i].y, self.boxcomplex.box[i].w,
                   self.boxcomplex.box[i].h,prob,labelname]
            print(box)
            boxes.append(box)

        return boxes, imgclasses

    def detect_image(self,size_img,image, boxes):

        thickness = (size_img.size[0] + size_img.size[1]) // 800
        for box in boxes:
            left = box[0]
            top = box[1]
            w = box[2]
            h = box[3]
            right = left + w
            bottom = top + h
            label = box[5]+' '+box[4]

            # left = max(0, np.floor(left + 0.5).astype('int32'))
            # bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            # right = min(image.size[0], np.floor(right + 0.5).astype('int32'))

            cv2.rectangle(image, (left, top), (right, bottom), self.colors[self.class_names.index(box[-1])], thickness//2)
            if (top > 10):
                cv2.putText(image, label, (left, top - 6), cv2.FONT_HERSHEY_COMPLEX_SMALL,thickness//3,  self.colors[self.class_names.index(box[-1])])
            else:
                cv2.putText(image, label, (left, top + 15), cv2.FONT_HERSHEY_COMPLEX_SMALL, thickness//3, self.colors[self.class_names.index(box[-1])])

        return image

    def scanAllImages(self,folderPath):
        extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in QImageReader.supportedImageFormats()]
        if '.tif' not in extensions:
            extensions.append('.tif')
        images = []
        for root, dirs, files in os.walk(folderPath):
            for file in files:
                if file.lower().endswith(tuple(extensions)):
                    relativePath = os.path.join(root, file)
                    path = str(os.path.abspath(relativePath))
                    images.append(path)
        images.sort(key=lambda x: x.lower())
        return images
